#!/usr/bin/env python
# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: donggu

import sys
import os
import subprocess


def get_env():
  env = {}
  env.update(os.environ)
  env['TZ'] = 'UTC'

  REPO_ROOT = os.path.dirname(os.path.realpath(__file__))

  # paths relative to repo root
  EXTRA_PYTHON_PATHS = [
      '/python',
      '/external/xunkemgmt/xunkemgmt_client/python',
  ]

  LD_LIBRARY_PATH = os.environ.get('LD_LIBRARY_PATH', '')

  if sys.version_info >= (3, 10, 0):
    EXTRA_PYTHON_PATHS.append('/../coin_binary/motion_lib/pyorio-0.0.1-py3.10-linux-x86_64.egg')
    EXTRA_PYTHON_PATHS.append('/../coin_binary/motion_lib/pypapyrus-0.0.1-py3.10-linux-x86_64.egg')
    EXTRA_PYTHON_PATHS.append('/../coin_binary/coin2-bazel-bin-3.10')
    EXTRA_PYTHON_PATHS.append('/../coin_binary/coin2-bazel-bin-3.10/external/apiset')
    LD_LIBRARY_PATH = f'../coin_binary/coin2-bazel-bin-3.10/_solib_k8:{LD_LIBRARY_PATH}'
  else:
    EXTRA_PYTHON_PATHS.append('/../coin_binary/motion_lib/pyorio-0.0.1-py3.8-linux-x86_64.egg')
    EXTRA_PYTHON_PATHS.append('/../coin_binary/coin2-bazel-bin/external/apiset')
    EXTRA_PYTHON_PATHS.append('/../coin_binary/coin2-bazel-bin')
    LD_LIBRARY_PATH = f'../coin_binary/coin2-bazel-bin/_solib_k8:{LD_LIBRARY_PATH}'

  PYTHONPATH = os.environ.get('PYTHONPATH', '')
  lst = [REPO_ROOT + p for p in EXTRA_PYTHON_PATHS]
  if PYTHONPATH: lst = lst + [PYTHONPATH]
  env['PYTHONPATH'] = ':'.join(lst)
  env['LD_LIBRARY_PATH'] = f'../coin_binary/motion_lib:{LD_LIBRARY_PATH}'

  return env

def file_to_module(path):
  path = os.path.normpath(path)
  assert os.path.isfile(path), path
  assert not os.path.isabs(path), path

  # a/b//c//d.py => [a,b,c,d]
  path, ext = os.path.splitext(path)
  assert ext == '.py', ext
  crumbs = filter(None, path.split('/'))
  # [a,b,c,d] => a.b.c.d
  module_path = '.'.join(crumbs)
  assert len(module_path) > 0

  return module_path


def main(argv):
  if argv[1] in ['-c', '-m', '-v']:
    os.execvpe('python', ['python'] + argv[1:], get_env())
  elif argv[1] == '-profile':
    os.execvpe('python', ['python', '-m', 'cProfile'] + argv[2:], get_env())
  elif argv[1] == '-gdb':
    env = get_env()
    env["PYTHONPATH"] = f'{os.getcwd()}:{env["PYTHONPATH"]}'
    subprocess.call("gdb --args python " + " ".join(argv[2:]), shell=True, env=env)
    return
  elif '--pdb' in argv:
    sys.stderr.write('[pyrunner] setup pdb.post_morten')
    pdb = True
    argv.remove('--pdb')
  else:
    pdb = False

  if len(argv) < 2:
    sys.stderr.write("No python file is given.\n")
    sys.stderr.write("Usage: ./pyrunner [--pdb] <python file>\n")
    return 1

  module_path = file_to_module(argv[1])

  if pdb:
    cmd = f"import pdb; import runpy; module_path='{module_path}'; exec('try: runpy.run_module(module_path, run_name=\\'__main__\\')\\nexcept: pdb.post_mortem()')"
    args = ['python', '-c', cmd] + argv[2:]
  else:
    args = ['python', '-m', module_path] + argv[2:]

  os.execvpe('python', args, get_env())

if __name__ == '__main__':
  sys.exit(main(sys.argv))
